from fastapi import FastAPI, File, UploadFile, Form
from fastapi.middleware.cors import CORSMiddleware
import tempfile
import os
from agent import create_qa_chain

app = FastAPI()

# 允许前端跨域
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.post("/ask")
async def ask_question(file: UploadFile = File(...), question: str = Form(...)):
    # 保存上传文件到临时目录
    with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1]) as tmp:
        tmp.write(await file.read())
        tmp_path = tmp.name

    try:
        qa_chain = create_qa_chain(tmp_path)
        answer = qa_chain.invoke({"query": question})
        return {"answer": answer["result"]}
    finally:
        os.unlink(tmp_path)  # 删除临时文件
